#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <netdb.h>
#include <malloc.h>
#include <getopt.h>
#include <arpa/inet.h>
#include <time.h>
#include <infiniband/verbs.h>

enum {
    PINGPONG_RECV_WRID = 1,
    PINGPONG_SEND_WRID = 2,
    MAX_QP             = 256,
};

static int page_size;
static int validate_buf;
static int use_odp;

struct pingpong_context {
    struct ibv_context  *context;
    struct ibv_comp_channel *channel;
    struct ibv_pd       *pd;
    struct ibv_mr       *mr;
    struct ibv_cq       *cq;
    struct ibv_srq      *srq;
    struct ibv_qp       *qp[MAX_QP];
    char            *buf;
    int          size;
    int          send_flags;
    int          num_qp;
    int          rx_depth;
    int          pending[MAX_QP];
    struct ibv_port_attr     portinfo;
};

struct pingpong_dest {
    int lid;
    int qpn;
    int psn;
    union ibv_gid gid;
};

void
wire_gid_to_gid(const char *wgid, union ibv_gid *gid)
{
    char tmp[9];
    __be32 v32;
    int i;
    uint32_t tmp_gid[4];

    for (tmp[8] = 0, i = 0; i < 4; ++i) {
        memcpy(tmp, wgid + i * 8, 8);
        sscanf(tmp, "%x", &v32);
        tmp_gid[i] = be32toh(v32);
    }
    memcpy(gid, tmp_gid, sizeof(*gid));
}

void
gid_to_wire_gid(const union ibv_gid *gid, char wgid[])
{
    uint32_t tmp_gid[4];
    int i;

    memcpy(tmp_gid, gid, sizeof(tmp_gid));
    for (i = 0; i < 4; ++i)
        sprintf(&wgid[i * 8], "%08x", htobe32(tmp_gid[i]));
}

static int
connect_ctx(struct pingpong_context *ctx, int port, enum ibv_mtu mtu,
    int sl, const struct pingpong_dest *my_dest,
    const struct pingpong_dest *dest, int sgid_idx)
{
    int i;
    /*修改QP状态为RTR*/
    for (i = 0; i < ctx->num_qp; ++i) {
        struct ibv_qp_attr attr = {
            .qp_state       = IBV_QPS_RTR,
            .path_mtu       = mtu,
            .dest_qp_num    = dest[i].qpn,
            .rq_psn         = dest[i].psn,
            .max_dest_rd_atomic = 1,
            .min_rnr_timer      = 12,
            .ah_attr        = {
                .is_global  = 0,
                .dlid       = dest[i].lid,
                .sl         = sl,
                .src_path_bits  = 0,
                .port_num   = port
            }
        };

        if (dest->gid.global.interface_id) {
            attr.ah_attr.is_global = 1;
            attr.ah_attr.grh.hop_limit = 1;
            attr.ah_attr.grh.dgid = dest->gid;
            attr.ah_attr.grh.sgid_index = sgid_idx;
        }
        if (ibv_modify_qp(ctx->qp[i], &attr,
                  IBV_QP_STATE              |
                  IBV_QP_AV                 |
                  IBV_QP_PATH_MTU           |
                  IBV_QP_DEST_QPN           |
                  IBV_QP_RQ_PSN             |
                  IBV_QP_MAX_DEST_RD_ATOMIC |
                  IBV_QP_MIN_RNR_TIMER)) {
            fprintf("Failed to modify QP[%d] to RTR\n", i);
            return 1;
        }
        /*修改状态为RTS*/
        attr.qp_state       = IBV_QPS_RTS;
        attr.timeout        = 14;
        attr.retry_cnt      = 7;
        attr.rnr_retry      = 7;
        attr.sq_psn     = my_dest[i].psn;
        attr.max_rd_atomic  = 1;
        if (ibv_modify_qp(ctx->qp[i], &attr,
                  IBV_QP_STATE              |
                  IBV_QP_TIMEOUT            |
                  IBV_QP_RETRY_CNT          |
                  IBV_QP_RNR_RETRY          |
                  IBV_QP_SQ_PSN             |
                  IBV_QP_MAX_QP_RD_ATOMIC)) {
            printf("Failed to modify QP[%d] to RTS\n", i);
            return 1;
        }
    }
    return 0;
}

static struct pingpong_context *
init_ctx(struct ibv_device *ib_dev, int ib_port, int num_qp)
{
    struct pingpong_context *ctx;
    int i, access_flags = IBV_ACCESS_LOCAL_WRITE;
    ctx = malloc(sizeof(*ctx));
    if(!ctx) {
        printf("ctx malloc error\n");
        return NULL;
    }
    memset(ctx, 0, sizeof(*ctx));
    ctx->size = 50;
    ctx->send_flags = IBV_SEND_SIGNALED;
    ctx->num_qp = num_qp;
    ctx->buf = memalign(page_size, size);
    if(!ctx->buf){
        printf("memalign error\n");
        goto clean_ctx;
    }
    memset(ctx->buf, 0, 50);

    /*打开设备*/
    ctx->context = ibv_open_device(ib_dev);
    if(!ctx->context){
        printf("Couldn't get context for %s\n",
            ibv_get_device_name(ib_dev));
        goto clean_buffer;
    }
    /*创建cc*/
    ctx->channel = ibv_create_comp_channel(ctx->context);
    if (!ctx->channel) {
        printf("Couldn't create completion channel\n");
        goto clean_device;
    }
    /*创建PD*/
    ctx->pd = ibv_alloc_pd(ctx->context);
    if (!ctx->pd) {
        printf("Couldn't allocate PD\n");
        goto clean_comp_channel;
    }
    /*创建MR*/
    ctx->mr = ibv_reg_mr(ctx->pd, ctx->buf, size, access_flags);
    if (!ctx->mr) {
        printf("Couldn't register MR\n");
        goto clean_pd;
    }
    /*创建CQ*/
    ctx->cq = ibv_create_cq(ctx->context, 10 + num_qp, NULL, ctx->channel, 0);
    if (!ctx->cq) {
        printf("Couldn't create CQ\n");
        goto clean_mr;
    }

    {
        struct ibv_srq_init_attr attr = {
            .attr = {
                .max_wr  = 10,
                .max_sge = 1
            }
        };
        /*创建SRQ*/
        ctx->srq = ibv_create_srq(ctx->pd, &attr);
        if (!ctx->srq)  {
            printf("Couldn't create SRQ\n");
            goto clean_cq;
        }
    }

    for (i = 0; i < num_qp; ++i) {
        struct ibv_qp_attr attr;
        struct ibv_qp_init_attr init_attr = {
            .send_cq = ctx->cq,
            .recv_cq = ctx->cq,
            .srq     = ctx->srq,/*SRP关联的QP*/
            .cap     = {
                .max_send_wr  = 1,
                .max_send_sge = 1,
            },
            .qp_type = IBV_QPT_RC/*QP类型RC*/
        };
        /*创建QP*/
        ctx->qp[i] = ibv_create_qp(ctx->pd, &init_attr);
        if (!ctx->qp[i])  {
            printf("Couldn't create QP[%d]\n", i);
            goto clean_qps;
        }
        ibv_query_qp(ctx->qp[i], &attr, IBV_QP_CAP, &init_attr);
        if (init_attr.cap.max_inline_data >= size) {
            ctx->send_flags |= IBV_SEND_INLINE;
        }
    }
    /*修改所有QP状态为INIT*/
    for (i = 0; i < num_qp; ++i) {
        struct ibv_qp_attr attr = {
            .qp_state        = IBV_QPS_INIT,
            .pkey_index      = 0,
            .port_num        = port,
            .qp_access_flags = 0
        };
        if (ibv_modify_qp(ctx->qp[i], &attr, IBV_QP_STATE | IBV_QP_PKEY_INDEX |
                   IBV_QP_PORT | IBV_QP_ACCESS_FLAGS)) {
            printf("Failed to modify QP[%d] to INIT\n", i);
            goto clean_qps_full;
        }
    }
    return ctx;
clean_qps_full:
    i = num_qp;
clean_qps:
    for (--i; i >= 0; --i)
        ibv_destroy_qp(ctx->qp[i]);
    ibv_destroy_srq(ctx->srq);
clean_cq:
    ibv_destroy_cq(ctx->cq);
clean_mr:
    ibv_dereg_mr(ctx->mr);
clean_pd:
    ibv_dealloc_pd(ctx->pd);
clean_comp_channel:
    if (ctx->channel)
        ibv_destroy_comp_channel(ctx->channel);
clean_device:
    ibv_close_device(ctx->context);
clean_buffer:
    free(ctx->buf);
clean_ctx:
    free(ctx);
    return NULL;
}

static int
close_ctx(struct pingpong_context *ctx, int num_qp)
{
    int i;
    /*销毁所有QP*/
    for (i = 0; i < num_qp; ++i) {
        if (ibv_destroy_qp(ctx->qp[i])) {
            printf("Couldn't destroy QP[%d]\n", i);
            return 1;
        }
    }
    /*销毁SRQ*/
    if (ibv_destroy_srq(ctx->srq)) {
        printf("Couldn't destroy SRQ\n");
        return 1;
    }
    /*销毁CQ*/
    if (ibv_destroy_cq(ctx->cq)) {
        printf("Couldn't destroy CQ\n");
        return 1;
    }
    /*销毁MR*/
    if (ibv_dereg_mr(ctx->mr)) {
        printf("Couldn't deregister MR\n");
        return 1;
    }
    /*销毁PD*/
    if (ibv_dealloc_pd(ctx->pd)) {
        printf("Couldn't deallocate PD\n");
        return 1;
    }
    /*销毁CC*/
    if (ctx->channel) {
        if (ibv_destroy_comp_channel(ctx->channel)) {
            printf("Couldn't destroy completion channel\n");
            return 1;
        }
    }
    /*关闭设备*/
    if (ibv_close_device(ctx->context)) {
        printf("Couldn't release context\n");
        return 1;
    }
    free(ctx->buf);
    free(ctx);
    return 0;
}

static int
post_recv(struct pingpong_context *ctx, int n)
{
    struct ibv_sge list = {
        .addr   = (uintptr_t) ctx->buf,
        .length = ctx->size,
        .lkey   = ctx->mr->lkey
    };
    struct ibv_recv_wr wr = {
        .wr_id      = PINGPONG_RECV_WRID,
        .sg_list    = &list,
        .num_sge    = 1,
    };
    struct ibv_recv_wr *bad_wr;
    int i;
    for (i = 0; i < n; ++i)
        if (ibv_post_srq_recv(ctx->srq, &wr, &bad_wr))
            break;
    return i;
}

static int
post_send(struct pingpong_context *ctx, int qp_index)
{
    struct ibv_sge list = {
        .addr   = (uintptr_t) ctx->buf,
        .length = ctx->size,
        .lkey   = ctx->mr->lkey
    };
    struct ibv_send_wr wr = {
        .wr_id      = PINGPONG_SEND_WRID,
        .sg_list    = &list,
        .num_sge    = 1,
        .opcode     = IBV_WR_SEND,
        .send_flags = ctx->send_flags,
    };
    struct ibv_send_wr *bad_wr;
    return ibv_post_send(ctx->qp[qp_index], &wr, &bad_wr);
}

static int
find_qp(int qpn, struct pingpong_context *ctx, int num_qp)
{
    int i;
    for (i = 0; i < num_qp; ++i)
        if (ctx->qp[i]->qp_num == qpn)
            return i;
    return -1;
}

static struct pingpong_dest *
server(struct pingpong_context *ctx,int ib_port, enum ibv_mtu mtu,
    int port, int sl,const struct pingpong_dest *my_dest, int sgid_idx)
{
    char gid[33];
    int n = 0;
    int sockfd = -1, connfd = -1;
    struct sockaddr_in serv_addr;
    struct pingpong_dest *rem_dest = NULL;
    char msg[sizeof "0000:000000:000000:00000000000000000000000000000000"];
    sockfd = socket(AF_INET, SOCK_STREAM, 0);
    if(sockfd < 0) {
        printf("socket error\n");
        return NULL;
    }

    memset(&serv_addr, 0, sizeof(serv_addr));
    serv_addr.sin_family = AF_INET;
    serv_addr.sin_addr.s_addr = htonl(INADDR_ANY);
    serv_addr.sin_port = htons(port); 

    bind(sockfd, (struct sockaddr*)&serv_addr, sizeof(serv_addr));
    if(listen(sockfd, 1) < 0) {
        printf("accept error\n");
        goto out;
    }

    connfd = accept(sockfd, (struct sockaddr*)NULL, NULL);
    if(connfd < 0) {
        printf("accept error\n");
        return NULL;
    }
    close(sockfd);

    rem_dest = malloc(MAX_QP * sizeof *rem_dest);
    if (!rem_dest) {
        printf("rem_dest malloc error\n");
        goto out;
    }

    for (i = 0; i < MAX_QP; ++i) {
        n = 0;
        while (n < sizeof msg) {
            /*循环读取发送过来的msg信息*/
            r = read(connfd, msg + n, sizeof msg - n);
            if (r < 0) {
                perror("server read");
                fprintf(stderr, "%d/%d: Couldn't read remote address [%d]\n",
                    n, (int) sizeof msg, i);
                goto out;
            }
            n += r;
        }
        /*将msg信息按照一定的格式输出到rem_dest的各个变量里*/
        sscanf(msg, "%x:%x:%x:%s", &rem_dest[i].lid, &rem_dest[i].qpn,
                &rem_dest[i].psn, gid);
        /*gid -> rem_dest[i].gid的转换*/
        wire_gid_to_gid(gid, &rem_dest[i].gid);
    }
    /*修改QP状态*/
    if (connect_ctx(ctx, ib_port, mtu, sl, my_dest, rem_dest, sgid_idx)) {
        printf("Couldn't connect to remote QP\n");
        free(rem_dest);
        rem_dest = NULL;
        goto out;
    }
    /*发送服务端的交换信息到客户端*/
    for (i = 0; i < MAX_QP; ++i) {
        /*my_dest[i].gid -> gid转换*/
        gid_to_wire_gid(&my_dest[i].gid, gid);
        sprintf(msg, "%04x:%06x:%06x:%s", my_dest[i].lid,
                    my_dest[i].qpn, my_dest[i].psn, gid);
        if (write(connfd, msg, sizeof msg) != sizeof msg) {
            printf("Couldn't send local address\n");
            free(rem_dest);
            rem_dest = NULL;
            goto out;
        }
    }
    /*读取客户端发送过来的信息，表示已经收到了服务端发送的交换信息*/
    if (read(connfd, msg, sizeof msg) != sizeof "done") {
        perror("client write");
        free(rem_dest);
        rem_dest = NULL;
        goto out;
    }
out:
    close(connfd);
    return rem_dest;
}

int main(int argc, char **argv)
{
    struct ibv_device      **dev_list;
    struct ibv_device       *ib_dev;
    struct pingpong_context *ctx;
    struct pingpong_dest     my_dest[MAX_QP];
    struct pingpong_dest    *rem_dest;
    char                    *servername = NULL;
    unsigned int             port = 18515;
    int                      ib_port = 1;
    int                      gidx = 1;
    char                     gid[33];
    int                      i;
    enum ibv_mtu             mtu = IBV_MTU_1024;
    unsigned int             num_qp = 16;

    page_size = sysconf(_SC_PAGESIZE);
    dev_list = ibv_get_device_list(NULL);
    if(!dev_list) {
        printf("ibv_get_device_list error\n");
        return 1;
    } else {
        ib_dev = dev_list[0];
        if(!ib_dev) {
            printf("IB device not found\n");
            return 1;
        }
    }

    ctx = init_ctx(ib_dev, ib_port, num_qp);
    if(!ctx) {
        printf("init_ctx error\n");
        return 1;
    }

    post_recv(ctx);

    if (ibv_req_notify_cq(ctx->cq, 0)) {
        printf("Couldn't request CQ notification\n");
        return 1;
    }

    memset(my_dest, 0, sizeof my_dest);
    if(ibv_query_port(ctx->context, ib_port, &ctx->portinfo)){
        printf("Couldn't get port info\n");
        return 1;
    }

    for(i = 0; i < num_qp; i++) {
        my_dest[i].qpn = ctx->qp[i]->qp_num;
        my_dest[i].psn = lrand48() & 0xffffff;
        my_dest[i].lid = ctx->portinfo.lid;
        /*获取gid的值，存储在my_dest[i].gid里*/
        if (ibv_query_gid(ctx->context, ib_port, gidx, &my_dest[i].gid)) {
            printf("Could not get local gid for gid index %d\n", gidx);
            return 1;
        }
        /*my_dest[i].gid -> gid(点分十进制表示)*/
        inet_ntop(AF_INET6, &my_dest[i].gid, gid, sizeof gid);
        printf("  local address:  LID 0x%04x, QPN 0x%06x, PSN 0x%06x, "
            "GID %s\n", my_dest[i].lid, my_dest[i].qpn, my_dest[i].psn, gid);
    }
    rem_dest = server(ctx, ib_port, mtu, sl, my_dest, rem_dest, gidx);
    if(!rem_dest) {
        return 1;
    }
    /*转成点分十进制的表示形式*/
    inet_ntop(AF_INET6, &rem_dest->gid, gid, sizeof gid);
    for (i = 0; i < num_qp; ++i) {
        inet_ntop(AF_INET6, &rem_dest[i].gid, gid, sizeof gid);
        printf("  remote address: LID 0x%04x, QPN 0x%06x, PSN 0x%06x, "
            "GID %s\n", rem_dest[i].lid, rem_dest[i].qpn,
            rem_dest[i].psn, gid);
    }

    int event_num = 0;
    while (event_num < 2) {
        struct ibv_cq *cq;
        struct ibv_wc wc;
        void *ex_ctx;

        ibv_get_cq_event(ctx->channel, &cq, &ex_ctx);
        ibv_ack_cq_events(cq, 1);
        event_num++;
        ibv_req_notify_cq(cq, 0);

        while(ibv_poll_cq(cq, 1, &wc)) {
            if(wc.status != IBV_WC_SUCCESS){
                printf("wc.status != IBV_WC_SUCCESS\n");
            }
            if(wc.opcode & IBV_WC_RECV) {
                printf("IBV_WC_RECV ok, %s\n", ctx->buf + 40);
                strncpy(ctx->buf + 40, "viscore server", sizeof "visocre server");
                if (post_send(ctx, rem_dest->qpn)) {
                    printf("Couldn't post send\n");
                    return 1;
                }
            }
            else if(wc.opcode == IBV_WC_SEND) {
                printf("IBV_WC_SEND ok\n");
            }
        }
    }
    if (close_ctx(ctx, num_qp)) {
        printf("close_ctx error\n");
        return 1;
    }
    ibv_free_device_list(dev_list);
    free(rem_dest);
    return 0;
}